'''
* This is the projet for brtc R&D  Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description MCP客户端实现模块，提供与MCP服务器交互的客户端功能
* @File: mcp_client.py
* @Time: 2025/01/27 14:30:00
* @All Rights Reserve By Brtc
'''

import json
import socket
from typing import Dict, Any, Optional, List

from .mcp_protocol import (
    MCPMessage,
    MCPResponse,
    OperationType,
    StatusCode,
    Context
)


class MCPClient:
    """MCP客户端类"""
    
    def __init__(self, host: str = "localhost", port: int = 8888):
        """
        * @Author Leon-liao
        * @Function: __init__
        * @Description //初始化MCP客户端
        * @Date :2025/01/27 14:30:00
        * @Param: host:服务器主机地址, port:服务器端口号
        * @return: 无
        """
        self.host = host
        self.port = port
        self.socket = None
    
    def connect(self):
        """
        * @Author Leon-liao
        * @Function: connect
        * @Description //连接到MCP服务器
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 无
        """
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect((self.host, self.port))
        print(f"已连接到MCP服务器 {self.host}:{self.port}")
    
    def disconnect(self):
        """
        * @Author Leon-liao
        * @Function: disconnect
        * @Description //断开与MCP服务器的连接
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 无
        """
        if self.socket:
            self.socket.close()
            self.socket = None
            print("已断开与MCP服务器的连接")
    
    def _send_request(self, message: MCPMessage) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: _send_request
        * @Description //发送请求并接收响应
        * @Date :2025/01/27 14:30:00
        * @Param: message:MCP消息对象
        * @return: 返回MCP响应对象
        """
        if not self.socket:
            raise ConnectionError("未连接到服务器，请先调用connect()")
        
        # 发送请求
        request_data = json.dumps(message.to_dict()).encode('utf-8')
        self.socket.sendall(request_data)
        
        # 接收响应
        response_data = self.socket.recv(4096)
        response_dict = json.loads(response_data.decode('utf-8'))
        response = MCPResponse.from_dict(response_dict)
        
        return response
    
    def create_context(
        self,
        context_id: Optional[str] = None,
        content: Any = None,
        metadata: Optional[Dict[str, Any]] = None,
        extensions: Optional[Dict[str, Any]] = None,
        scope: Optional[str] = None
    ) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: create_context
        * @Description //创建新的上下文
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID, content:上下文内容, metadata:元数据, 
                 extensions:扩展字段, scope:作用域
        * @return: 返回MCP响应对象
        """
        message = MCPMessage(
            operation=OperationType.CREATE,
            context_id=context_id,
            scope=scope,
            payload={
                "content": content,
                "metadata": metadata or {},
                "extensions": extensions or {}
            }
        )
        return self._send_request(message)
    
    def read_context(self, context_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: read_context
        * @Description //读取指定ID的上下文
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID
        * @return: 返回MCP响应对象
        """
        message = MCPMessage(
            operation=OperationType.READ,
            context_id=context_id
        )
        return self._send_request(message)
    
    def update_context(
        self,
        context_id: str,
        content: Optional[Any] = None,
        metadata: Optional[Dict[str, Any]] = None,
        extensions: Optional[Dict[str, Any]] = None
    ) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: update_context
        * @Description //更新指定ID的上下文
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID, content:新的上下文内容, 
                 metadata:新的元数据, extensions:新的扩展字段
        * @return: 返回MCP响应对象
        """
        payload = {}
        if content is not None:
            payload["content"] = content
        if metadata is not None:
            payload["metadata"] = metadata
        if extensions is not None:
            payload["extensions"] = extensions
        
        message = MCPMessage(
            operation=OperationType.UPDATE,
            context_id=context_id,
            payload=payload
        )
        return self._send_request(message)
    
    def delete_context(self, context_id: str) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: delete_context
        * @Description //删除指定ID的上下文
        * @Date :2025/01/27 14:30:00
        * @Param: context_id:上下文ID
        * @return: 返回MCP响应对象
        """
        message = MCPMessage(
            operation=OperationType.DELETE,
            context_id=context_id
        )
        return self._send_request(message)
    
    def search_contexts(
        self,
        labels: Optional[List[str]] = None,
        scope: Optional[str] = None,
        keyword: Optional[str] = None
    ) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: search_contexts
        * @Description //搜索上下文
        * @Date :2025/01/27 14:30:00
        * @Param: labels:标签列表, scope:作用域, keyword:关键词
        * @return: 返回MCP响应对象
        """
        criteria = {}
        if labels:
            criteria["labels"] = labels
        if scope:
            criteria["scope"] = scope
        if keyword:
            criteria["keyword"] = keyword
        
        message = MCPMessage(
            operation=OperationType.SEARCH,
            payload={"criteria": criteria}
        )
        return self._send_request(message)
    
    def copy_context(
        self,
        source_id: str,
        target_id: str
    ) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: copy_context
        * @Description //复制上下文
        * @Date :2025/01/27 14:30:00
        * @Param: source_id:源上下文ID, target_id:目标上下文ID
        * @return: 返回MCP响应对象
        """
        message = MCPMessage(
            operation=OperationType.COPY,
            context_id=source_id,
            payload={"target_id": target_id}
        )
        return self._send_request(message)
    
    def list_tools(self) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: list_tools
        * @Description //列出服务器中所有可用的工具
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回MCP响应对象，包含工具列表
        """
        message = MCPMessage(operation=OperationType.LIST_TOOLS)
        return self._send_request(message)
    
    def call_tool(
        self,
        tool_name: str,
        arguments: Optional[Dict[str, Any]] = None
    ) -> MCPResponse:
        """
        * @Author Leon-liao
        * @Function: call_tool
        * @Description //调用指定的工具
        * @Date :2025/01/27 14:30:00
        * @Param: tool_name:工具名称, arguments:工具参数字典
        * @return: 返回MCP响应对象，包含工具执行结果
        """
        message = MCPMessage(
            operation=OperationType.CALL_TOOL,
            payload={
                "tool_name": tool_name,
                "arguments": arguments or {}
            }
        )
        return self._send_request(message)
    
    def __enter__(self):
        """
        * @Author Leon-liao
        * @Function: __enter__
        * @Description //上下文管理器入口，支持with语句
        * @Date :2025/01/27 14:30:00
        * @Param: 无
        * @return: 返回自身对象
        """
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        * @Author Leon-liao
        * @Function: __exit__
        * @Description //上下文管理器出口，自动断开连接
        * @Date :2025/01/27 14:30:00
        * @Param: exc_type:异常类型, exc_val:异常值, exc_tb:异常追踪
        * @return: 返回False表示不抑制异常
        """
        self.disconnect()
        return False

